# -*- coding: utf-8 -*-
# time: 2025/4/30 15:56
# file: 数学计算_Agent.py
# author: hanson
from langchain_community.chat_models import ChatOllama
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain.tools import Tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

# 1. 定义数学计算工具
def math_calculator(expression: str) -> str:
    """计算数学表达式，如 '2+2' 或 'sqrt(9)'"""
    try:
        result = eval(expression)  # 注意：实际生产环境应使用更安全的计算方式
        return str(result)
    except Exception as e:
        return f"计算错误: {e}"

math_tool = Tool.from_function(
    name="math_tool",
    description="用于计算数学表达式",
    func=math_calculator
)

# 2. 创建 LLM
llm = ChatOllama(model="qwen2.5:1.5b", temperature=0)

# 3. 定义正确的提示模板（包含agent_scratchpad）
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个专业的数学助手，能够准确计算各种数学表达式。"),
    ("user", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad")  # 这是关键部分
])

# 4. 创建 Agent
agent = create_tool_calling_agent(llm, [math_tool], prompt)

# 5. 创建执行器
agent_executor = AgentExecutor(agent=agent, tools=[math_tool], verbose=True)

# 6. 使用 Agent
result = agent_executor.invoke({"input": "计算3的4次方加上10"})
print(result["output"])